#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include "gen_riscv.h"
#include "ir.h"
#include "ast.h"

// --- 新的数据结构：用于映射 IR 变量到栈偏移 ---
typedef struct VarMap {
    char name[32];
    int offset;
    struct VarMap *next;
} VarMap;

/* --- 修复：为 is_number 添加前向声明 --- */
static int is_number(const char *s);

// --- 辅助函数：加载/存储 ---
static void load_var(FILE *out, const char* physical_reg, const char* ir_var, VarMap* var_map) {
    if (is_number(ir_var)) {
        fprintf(out, "  li %s, %s\n", physical_reg, ir_var);
        return;
    }
    for (VarMap *v = var_map; v; v = v->next) {
        if (strcmp(v->name, ir_var) == 0) {
            fprintf(out, "  lw %s, %d(sp)\n", physical_reg, v->offset);
            return;
        }
    }
    fprintf(out, "  # WARNING: Var '%s' not found in stack map, loading 0\n", ir_var);
    fprintf(out, "  li %s, 0\n", physical_reg);
}

static void store_var(FILE *out, const char* physical_reg, const char* ir_var, VarMap* var_map) {
    for (VarMap *v = var_map; v; v = v->next) {
        if (strcmp(v->name, ir_var) == 0) {
            fprintf(out, "  sw %s, %d(sp)\n", physical_reg, v->offset);
            return;
        }
    }
    fprintf(out, "  # WARNING: Var '%s' not found in stack map, cannot store\n", ir_var);
}

/* is_number 的定义，放到 load_var 之后 */
static int is_number(const char *s) {
    if (!s || *s == '\0') return 0;
    char *endptr;
    strtol(s, &endptr, 10);
    return *endptr == '\0';
}

// --- 主生成函数 ---
void gen_riscv(FILE *out) {
    fprintf(out, ".text\n");
    IR *ir = ir_head;

    while (ir) {
        if (ir->op == IR_FUNC_BEGIN) {
            const char* func_name = ir->a;
            
            VarMap *var_map = NULL;
            int current_offset = 0;
            current_offset += 4; // For ra
            
            IR* p = ir->next;
            while (p && p->op != IR_FUNC_END) {
                char* vars_to_map[] = {p->a, p->b, p->c};
                for (int i = 0; i < 3; ++i) {
                    char* var = vars_to_map[i];
                    if (var && var[0] != '\0' && !is_number(var)) {
                        int found = 0;
                        for (VarMap *v = var_map; v; v = v->next) {
                            if (strcmp(v->name, var) == 0) {
                                found = 1;
                                break;
                            }
                        }
                        if (!found) {
                            VarMap *new_var = (VarMap*)malloc(sizeof(VarMap));
                            strncpy(new_var->name, var, sizeof(new_var->name)-1);
                            new_var->offset = current_offset;
                            current_offset += 4;
                            new_var->next = var_map;
                            var_map = new_var;
                        }
                    }
                }
                p = p->next;
            }
            
            int stack_size = (current_offset + 15) / 16 * 16;

            fprintf(out, "\n.globl %s\n", func_name);
            fprintf(out, "%s:\n", func_name);
            fprintf(out, "  addi sp, sp, -%d\n", stack_size);
            fprintf(out, "  sw ra, %d(sp)\n", stack_size - 4);

            int arg_count_for_call = 0;
            p = ir->next;
            while (p && p->op != IR_FUNC_END) {
                switch (p->op) {
                    case IR_LABEL:
                        fprintf(out, "%s:\n", p->a);
                        break;
                    case IR_GOTO:
                        fprintf(out, "  j %s\n", p->a);
                        break;
                    case IR_IFZ:
                        load_var(out, "t0", p->a, var_map);
                        fprintf(out, "  beqz t0, %s\n", p->b);
                        break;
                    case IR_LOADI:
                        fprintf(out, "  li t0, %d\n", p->extra);
                        store_var(out, "t0", p->a, var_map);
                        break;
                    case IR_MOVE:
                        if (strncmp(p->b, "param", 5) == 0) {
                            int param_idx = atoi(p->b + 5);
                            if (param_idx < 8) {
                                fprintf(out, "  mv t0, a%d\n", param_idx);
                                store_var(out, "t0", p->a, var_map);
                            }
                        } else {
                            load_var(out, "t0", p->b, var_map);
                            store_var(out, "t0", p->a, var_map);
                        }
                        break;
                    case IR_BIN:
                        load_var(out, "t1", p->b, var_map);
                        load_var(out, "t2", p->c, var_map);
                        
                        switch((Operator)p->extra) {
                            case OP_PLUS:  fprintf(out, "  add t0, t1, t2\n"); break;
                            case OP_MINUS: fprintf(out, "  sub t0, t1, t2\n"); break;
                            case OP_MUL:   fprintf(out, "  mul t0, t1, t2\n"); break;
                            case OP_DIV:   fprintf(out, "  div t0, t1, t2\n"); break;
                            case OP_MOD:   fprintf(out, "  rem t0, t1, t2\n"); break;
                            case OP_LT:    fprintf(out, "  slt t0, t1, t2\n"); break;
                            case OP_GT:    fprintf(out, "  sgt t0, t1, t2\n"); break;
                            case OP_LE:    fprintf(out, "  sgt t0, t1, t2\n  xori t0, t0, 1\n"); break;
                            case OP_GE:    fprintf(out, "  slt t0, t1, t2\n  xori t0, t0, 1\n"); break;
                            case OP_EQ:    fprintf(out, "  sub t3, t1, t2\n  seqz t0, t3\n"); break;
                            case OP_NE:    fprintf(out, "  sub t3, t1, t2\n  snez t0, t3\n"); break;
                            default: break;
                        }
                        store_var(out, "t0", p->a, var_map);
                        break;
                    case IR_RET:
                        if (p->a[0] != '\0') {
                            load_var(out, "a0", p->a, var_map);
                        } else {
                            fprintf(out, "  li a0, 0\n");
                        }
                        fprintf(out, "  j %s_epilogue\n", func_name);
                        break;
                    case IR_ARG:
                        if (arg_count_for_call < 8) {
                            char reg_name[4];
                            sprintf(reg_name, "a%d", arg_count_for_call);
                            /* --- 修复：使用正确的参数调用 load_var --- */
                            load_var(out, reg_name, p->a, var_map);
                        }
                        arg_count_for_call++;
                        break;
                    case IR_CALL:
                        fprintf(out, "  call %s\n", p->b);
                        store_var(out, "a0", p->a, var_map);
                        arg_count_for_call = 0;
                        break;
                    default: break;
                }
                p = p->next;
            }
            
            fprintf(out, "%s_epilogue:\n", func_name);
            fprintf(out, "  lw ra, %d(sp)\n", stack_size - 4);
            fprintf(out, "  addi sp, sp, %d\n", stack_size);
            fprintf(out, "  ret\n");

            while(var_map) {
                VarMap* temp = var_map;
                var_map = var_map->next;
                free(temp);
            }

            ir = p;
        }
        
        if (ir) ir = ir->next;
    }
}